import chisel3._
import chisel3.util._
import Control._

class MEU(xlen: Int)extends Module {
  val io = IO(new Bundle {
    val in = Flipped(Decoupled(new MessageEM(xlen)))
    val out = Decoupled(new MessageMW(xlen))

    val meu_valid = Output(Bool())
    val addr      = Output(UInt(xlen.W))
    val st_type   = Output(UInt(2.W))
    val ld_type   = Output(UInt(3.W))
    val wr_data   = Output(UInt(xlen.W))
    val rd_data   = Input(UInt(xlen.W))
    val mem_valid = Input(Bool())

    // output br_taken to IFU and IDU
    val br_taken = Output(Bool())
    val Alu_out  = Output(UInt(xlen.W))
  })

  // pipeline registers (exu => meu)
  val meu_inst     = RegInit(0.U(xlen.W))
  val meu_pc       = RegInit(0.U(xlen.W))
  val meu_PC_sel   = RegInit(0.U(2.W))
  val meu_st_type  = RegInit(0.U(2.W))
  val meu_ld_type  = RegInit(0.U(3.W))
  val meu_wb_sel   = RegInit(0.U(2.W))
  val meu_csr_cmd  = RegInit(0.U(3.W))
  val meu_wb_en    = RegInit(0.B)
  val meu_br_taken = RegInit(0.B) // NOTE: for simulation to identify next pc
  val meu_Alu_out  = RegInit(0.U(xlen.W))
  val meu_wb_addr  = RegInit(0.U(5.W))
  val meu_rs2_data = RegInit(0.U(xlen.W))

  val in_valid  = RegInit(0.B)  // valid for input
  val out_valid = RegInit(0.B)  // valid for output
  val ready     = RegInit(1.B)  // ready for input
  val rd_data   = RegInit(0.U(xlen.W))

  val br_taken = RegInit(0.B) // for temporarily saving the input bt_taken

  meu_inst     := Mux(io.in.valid && io.in.ready, io.in.bits.inst,     meu_inst)
  meu_pc       := Mux(io.in.valid && io.in.ready, io.in.bits.pc,       meu_pc)
  meu_PC_sel   := Mux(io.in.valid && io.in.ready, io.in.bits.PC_sel,   meu_PC_sel)
  meu_st_type  := Mux(io.in.valid && io.in.ready, io.in.bits.st_type,  meu_st_type)
  meu_ld_type  := Mux(io.in.valid && io.in.ready, io.in.bits.ld_type,  meu_ld_type)
  meu_wb_sel   := Mux(io.in.valid && io.in.ready, io.in.bits.wb_sel,   meu_wb_sel)
  meu_csr_cmd  := Mux(io.in.valid && io.in.ready, io.in.bits.csr_cmd,  meu_csr_cmd)
  meu_wb_en    := Mux(io.in.valid && io.in.ready, io.in.bits.wb_en,    meu_wb_en)
  meu_br_taken := Mux(io.in.valid && io.in.ready, io.in.bits.br_taken, meu_br_taken) // NOTE: for simulation to identify next pc
  meu_Alu_out  := Mux(io.in.valid && io.in.ready, io.in.bits.Alu_out,  meu_Alu_out)
  meu_wb_addr  := Mux(io.in.valid && io.in.ready, io.in.bits.wb_addr,  meu_wb_addr)
  meu_rs2_data := Mux(io.in.valid && io.in.ready, io.in.bits.rs2_data, meu_rs2_data)

  br_taken := Mux(br_taken, 0.B, Mux(io.in.valid && io.in.ready, io.in.bits.br_taken, br_taken))

  rd_data := Mux(in_valid & io.mem_valid, io.rd_data, rd_data)

  io.out.bits.inst    := meu_inst
  io.out.bits.pc      := meu_pc

  io.out.bits.PC_sel   := meu_PC_sel
  io.out.bits.wb_sel   := meu_wb_sel
  io.out.bits.csr_cmd  := meu_csr_cmd
  io.out.bits.wb_en    := meu_wb_en
  io.out.bits.br_taken := meu_br_taken // NOTE: for simulation to identify next pc

  io.out.bits.Alu_out := meu_Alu_out
  io.out.bits.wb_addr := meu_wb_addr
  io.out.bits.rd_data := io.rd_data

  io.meu_valid := in_valid
  io.addr      := meu_Alu_out
  io.st_type   := meu_st_type
  io.ld_type   := meu_ld_type
  io.wr_data   := meu_rs2_data

  io.br_taken  := br_taken
  io.Alu_out   := meu_Alu_out

  // in_valid  := Mux(io.in.valid & io.in.ready,   1.B, Mux(in_valid & io.mem_valid, 0.B, in_valid))
  in_valid  := Mux(io.in.valid & io.in.ready,   1.B, Mux(io.out.valid & io.out.ready, 0.B, in_valid))
  out_valid := Mux(io.out.valid & io.out.ready, 0.B, Mux(in_valid & io.mem_valid, 1.B, out_valid))
  io.out.valid := out_valid

  ready := Mux(io.in.valid & io.in.ready, 0.B, Mux(io.out.valid & io.out.ready, 1.B, ready))
  io.in.ready  := ready
}
